#include "sf_rc.h"

#include <random>
#include <cassert>

using namespace skyfire;
using namespace arma;
using namespace std;

sf_rc::sf_rc(const std::string& ps) 
{
    init_u__();
    string tmp_ps = "0.";
    for(auto &p:ps)
        tmp_ps += to_string(static_cast<int>(p));
    password__ = tmp_ps;
    for(int i = 0;i<SF_RC_PASSWORD_INIT_TIMES;++i)
        password__ = (password__ * u__ * (1 - password__));
}

void sf_rc::init_u__()
{
    string str(SF_RC_PASSWORD_BIT_COUNT,'9');
    str[0]='3';
    str[1]='.';
    u__ = str;
}

std::vector<sf_float> sf_rc::make_logistic_col__(int len)
{
    std::vector<sf_float> ret;
    for(int i=0;i<len;++i)
    {
        password__ = (password__ * u__ * (1 - password__));
        ret.push_back(password__);
    }
    return ret;
}

Col<sf_byte> sf_rc::map_to_n_byte__(const std::vector<sf_float> &data,sf_byte n)
{
    Col<sf_byte > ret(data.size());
    for(int i = 0;i<data.size();++i)
        ret(i)= sf_float(data[i]*n).get_si();
    return ret;
}

Col<int> sf_rc::map_to_n_int__(const std::vector<sf_float> &data, int n)
{
    Col<int > ret(data.size());
    for(int i = 0;i<data.size();++i)
        ret(i)= sf_float(data[i]*n).get_si();
    return ret;
}

Col<sf_byte> sf_rc::make_random_list__(int len)
{
    Col<sf_byte > ret(len);
    random_device rd;
    default_random_engine e(rd());
    uniform_int_distribution<sf_byte > ed(0, 255);
    for(int i = 0;i<len;++i)
        ret(i) = ed(e);
    return ret;
}

Mat<sf_byte> sf_rc::make_chaos_mat__()
{
    Mat<sf_byte> ret(16,16);
    for(int i=0;i<16;++i)
        ret.col(i) = map_to_n_byte__(make_logistic_col__(16), 255);
    return ret;
}


sf_byte_array sf_rc::enc_block__(const sf_byte_array& data)
{
    assert(data.size()<=256);
    if(data.size()==0)
        return sf_byte_array();
    Mat<sf_byte> data_block(data);

    // 一、数据及密钥准备

    // 1. 记录原始数据长度（1~256）
    sf_byte size = static_cast<sf_byte>(data.size());

    // 2. 将原始数据reshape为16*16的矩阵
    data_block.reshape(16,16);

    // 3. 生成随机行列密钥向量（16）
    Row<sf_byte> rand_row_key = make_random_list__(16).t();
    Col<sf_byte> rand_col_key = make_random_list__(16);

    // 4. 根据随机行列密钥向量生成16*16随机密钥矩阵
    Mat<sf_byte> rand_mat_key = rand_col_key * rand_row_key;

    // 5. 生成16*16混沌矩阵
    Mat<sf_byte> chaos_mat_key = make_chaos_mat__();

    // 6. 生成长度为289的混沌置换列向量
    Col<int> chaos_swap_key = map_to_n_int__(make_logistic_col__(17*17),17*17);
    
    // 二、加密过程

    // 1. 使用混沌矩阵和随机密钥矩阵对源数据进行异或加密
    for(int i = 0;i<16;++i)
        for(int j = 0;j<16;++j)
            data_block(i,j) = data_block(i,j) ^ chaos_mat_key(i,j) ^ rand_mat_key(i,j);

    // 2. 附加随机列密钥向量追加至加密数据列尾
    data_block.insert_cols(16, rand_col_key);

    // 3. 将数据大小追加至随机行密钥向量列尾
    Col<sf_byte> tmp_col(1);
    tmp_col(0) = size;
    rand_row_key.insert_cols(16,tmp_col);

    // 4. 将处理后的随机行密钥向量追加至加密数据行尾
    data_block.insert_rows(16, rand_row_key);
    
    // 5. 将加密数据reshape为1*289的行向量
    data_block.reshape(1,17*17);

    // 6. 使用混沌置换列向量交换密文中的元素，让随机密钥隐藏至密文中
    for(int i=0;i<17*17;++i)
        data_block.swap_cols(i, chaos_swap_key(i));

    // 7. 加密完成
    return {data_block.begin(), data_block.end()};
}


sf_byte_array sf_rc::dec_block__(const sf_byte_array& data)
{
    assert(data.size()==289);
    Mat<sf_byte> data_block(data);
    data_block.reshape(1 , 17 * 17);

    Row<sf_byte> rand_row_key;
    Col<sf_byte> rand_col_key;
    Mat<sf_byte> chaos_mat_key = make_chaos_mat__();
    Col<int> chaos_swap_key = map_to_n_int__(make_logistic_col__(17*17),17*17);
    
    for(int i=17*17-1;i>=0;--i)
        data_block.swap_cols(i, chaos_swap_key(i));

    data_block.reshape(17,17);
    rand_row_key = data_block.row(16);
    data_block.shed_row(16);
    rand_col_key = data_block.col(16);
    data_block.shed_col(16);
    sf_byte size = rand_row_key(16);
    rand_row_key.shed_col(16);

    Mat<sf_byte> rand_mat_key = rand_col_key * rand_row_key;

    for(int i = 0;i<16;++i)
        for(int j = 0;j<16;++j)
            data_block(i,j) = data_block(i,j) ^ chaos_mat_key(i,j) ^ rand_mat_key(i,j);
    data_block.reshape(1,16*16);
    if(size == 0)
        return {data_block.begin(),data_block.end()};
    return {data_block.begin(), data_block.begin()+size};
}


sf_byte_array sf_rc::enc(const sf_byte_array& data)
{
    int group_num = data.size() / 256;
    sf_byte_array ret;
    for(int i = 0; i<group_num;++i)
    {
        ret += enc_block__({data.begin()+i*256,data.begin()+(i+1)*256});
    }
    ret += enc_block__({data.begin()+group_num*256, data.end()});
    return ret;
} 

sf_byte_array sf_rc::dec(const sf_byte_array& data)
{
    assert(data.size()%289==0);
    int group_num = data.size() / 289;
    sf_byte_array ret;
    for(int i = 0; i<group_num;++i)
    {
        ret += dec_block__({data.begin()+i*289,data.begin()+(i+1)*289});
    }
    return ret;
}

bool sf_rc::enc(istream *in, ostream * out)
{
    if(!*in || !*out)
        return false;
    sf_byte_array data;
    while(*in)
    {
        data.resize(256);
        in->read(reinterpret_cast<char*>(data.data()), 256);
        data.resize(in->gcount());
        if(!out->write(reinterpret_cast<char*>(enc(data).data()), 289))
            return false;
    }
    return true;
}

bool sf_rc::dec(istream *in, ostream * out)
{
    if(!*in || !*out)
        return false;
    sf_byte_array data;
    sf_byte_array dec_data;
    while(*in)
    {
        data.resize(289);
        in->read(reinterpret_cast<char*>(data.data()), 289);
        data.resize(in->gcount());
        if(data.size() == 0)
            break;
        if(data.size() != 289)
            return false;
        dec_data = dec(data);
        if(!out->write(reinterpret_cast<char*>(dec_data.data()), dec_data.size()))
            return false;
    }
    return true;
}